# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import logging
import sys
from contextlib import contextmanager

import kafka
from absl import app, flags

from coin.support.proto_log.logic.log_info import (
    OgInfo, StratInfo, SubscriptionInfo)
from coin.util.queue.config import KafkaConfig
from coin.util.queue.constants import TopicType


@contextmanager
def kafka_logger_level(level):
  kafka_logger = logging.getLogger('kafka')
  org_level = kafka_logger.getEffectiveLevel()
  kafka_logger.setLevel(level)
  yield
  kafka_logger.setLevel(org_level)


def create_topics(topics, kafka_config):
  with kafka_logger_level(logging.FATAL):
    client = kafka.KafkaClient(kafka_config.kafka_servers)
    try:
      topics = kafka_config.gen_topics(topics)
      for topic in topics:
        client.ensure_topic_exists(topic)
        sys.stderr.write('Topic created: %s\n' % topic)
    finally:
      client.close()


def delete_topics(topics, kafka_config):
  with kafka_logger_level(logging.WARNING):
    admin_client = kafka.admin.KafkaAdminClient(bootstrap_servers=kafka_config.kafka_servers)
    try:
      topics = kafka_config.gen_topics(topics)
      admin_client.delete_topics(topics)
      sys.stderr.write('Topic deleted: %s\n' % ', '.join(topics))
    except kafka.errors.UnknownTopicOrPartitionError:
      sys.stderr.write('Unknown topic(s).\n')
    finally:
      admin_client.close()


def list_topics(kafka_config):
  with kafka_logger_level(logging.WARNING):
    logging.getLogger('kafka').setLevel(logging.WARNING)
    client = kafka.KafkaClient(kafka_config.kafka_servers)
    try:
      topics = client.topics
      for topic in sorted(topics):
        print(topic)
    finally:
      client.close()


def generate_kafka_topic(topic_type, topic_info):
  assert isinstance(topic_type, TopicType)
  topic = None
  if topic_type in (TopicType.TICKER,):
    assert isinstance(topic_info, SubscriptionInfo)
    sub_key = f'{topic_info.market_type}.{topic_info.exchange}.{topic_info.api_version}'
    topic = f'{topic_type.name}_{topic_info.machine}_{sub_key}'
  elif topic_type in (TopicType.ACCOUNT, TopicType.ACCOUNTESTIMATE):
    assert isinstance(topic_info, OgInfo)
    acct_key = f'{topic_info.market_type}.{topic_info.exchange}.{topic_info.owner}'
    topic = f'{topic_type.name}_{topic_info.machine}_{acct_key}'
  elif topic_type in (TopicType.STRAT, TopicType.TELEMETRY):
    assert isinstance(topic_info, StratInfo)
    topic = f'{topic_type.name}_{topic_info.machine}_{topic_info.strategy_name}'
  else:
    raise ValueError('Unknow topic type: %s' % topic_type.name)
  return topic


def parse_kafka_topic(topic):
  topic_type, machine, info = topic.split('_', 2)
  topic_type = TopicType[topic_type]
  if topic_type in (TopicType.TICKER,):
    market_type, exchange, api_version = info.split('.')
    topic_info = SubscriptionInfo(
        market_type=market_type, exchange=exchange, api_version=api_version,
        machine=machine, trading_date=None)
  elif topic_type in (TopicType.ACCOUNT, TopicType.ACCOUNTESTIMATE):
    market_type, exchange, owner = info.split('.')
    topic_info = OgInfo(market_type=market_type, exchange=exchange, owner=owner,
                        machine=machine, trading_date=None)
  elif topic_type in (TopicType.STRAT, TopicType.TELEMETRY,):
    topic_info = StratInfo(strategy_name=info, machine=machine,
                           trading_date=None)
  else:
    raise ValueError('Unknow topic type: %s' % topic_type.name)
  return topic_type, topic_info


def main(argv):
  FLAGS = flags.FLAGS
  kafka_config = KafkaConfig.from_cmd_config(FLAGS.kafka_config_filename)

  if FLAGS.create:
    create_topics(FLAGS.create.split(','), kafka_config)

  if FLAGS.delete:
    delete_topics(FLAGS.delete.split(','), kafka_config)

  if FLAGS.list or (not FLAGS.create and not FLAGS.delete):
    list_topics(kafka_config)


if __name__ == "__main__":
  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')

  flags.DEFINE_bool('list', None, 'List all topics')

  flags.DEFINE_string('create', None, 'Topic to be created.')

  flags.DEFINE_string('delete', None, 'Topic to be deleted.')

  logging.getLogger('kafka').setLevel(logging.WARNING)
  app.run(main)
